from pylab import *
from scipy import *
from scipy import optimize
from matplotlib import pyplot
import math

dir = '/home/jcfernandez/resultados/eps1/u100/AMPL/W2.0/STRESS/'
name='strE1U100W2A200.dat'
omega=2.0
repair = False
plotgraph =  True
if repair:
    valMax=50
    datos=[]
    mmm=0.;
    for line in file(dir+name):
        line=line.replace('\n','')
        val=line.split(' ')
        res=[]
        data=[]
        for v in val:
            if v!='':
                res.append(v)
        if res[1].replace('-','')!='nan':
            if abs(float(res[1]))<valMax:
                datos.append([float(res[0]),float(res[1])])
                if mmm<float(res[1]):
                    mmm=float(res[1])
    x=[]
    y=[]
    f = open(dir+'neo-'+name,'w')
    for dt in datos:
        f.write(str(dt[0])+' '+str(dt[1])+'\n')
        x.append(dt[0])
        y.append(dt[1])
    f.close()
    
    pyplot.plot(x,y,'ro')
    pyplot.show()


if not repair:
    
    offset=2.0
    x=[]
    y=[]
    for line in file(dir+name):
        line=line.replace('\n','')
        val=line.split(' ')
        dt=[]
        for v in val:
            if v!='':
                dt.append(float(v))
        if dt[0]>offset:
            x.append(dt[0])
            y.append(dt[1])
    x = array(x)
    y = array(y)
    #num_points = 150
    #Tx = linspace(5., 8., num_points)
    #Ty = Tx
    #tX = 11.86*cos(2*pi/0.81*Tx-1.32) + 0.64*Tx+4*((0.5-rand(num_points))*exp(2*rand(num_points)**2))
    #tY = -32.14*cos(2*pi/0.8*Ty-1.94) + 0.15*Ty+7*((0.5-rand(num_points))*exp(2*rand(num_points)**2))
    #
    ## Fit the first set
    fitfunc = lambda p, x: p[0]*sin(omega*x+p[1])# Target function
    errfunc = lambda p, x, y: fitfunc(p, x) - y # Distance to the target function
    p0 = [4., 1.0] # Initial guess for the parameters
    #p1, success = optimize.leastsq(errfunc, p0[:], args=(x, y))
    out = optimize.leastsq(errfunc, p0[:], args=(x, y),full_output=1)
    p1 = out[0]
    covar = out[1]
    xErr = sqrt( covar[0][0] )
    yErr = sqrt( covar[1][1] )
    print "Tau0: ",p1[0],xErr
    print "delta: ",p1[1],yErr
    print omega,p1[0],p1[1]
    time = linspace(x.min(), x.max(), 1000)
    plot(x, y, "ro", time, fitfunc(p1, time), "b-") # Plot of the data and the fit
    #plot(x, y,"ro") 
    show()